#pragma once

#include <vector>
#include <list>
#include <limits>
#include "types.h"
#include <typelib/types.h>
#include <imglib/color.h>
#include <imglib/image.h>

namespace refine {

typedef signed char schar;
typedef unsigned short ushort;
typedef unsigned int uint;

typedef point<schar> spoint; ///< relative coordinates of a neighbour pixel
typedef rect<schar> srect; 
typedef std::pair<ushort, spoint> cdist_spoint_pair; ///< memory saving is significant here
typedef std::vector<cdist_spoint_pair> neighbours;
typedef img::BGR<int> icolor;

/// The object stores coordinates of neighbours for some points,
/// which are similar to given colors, in sorted order.
/// The sorting is very time-consuming, that is why caching is required.
template <class R>
class sorted_neighbours {
public:
	sorted_neighbours(const R & im_, uint whs_, information consider_);

	const R & image() const
		{ return im; }
	uint size() const ///< the number of elements in the cache (each element correspond to a pixel)
		{ return pool.size(); }
	uint scans() const ///< the total number of neighbours scans made during lifetime of the object
		{ return total_scans; }
	uint sorts() const ///< the total number of partial sorts made during lifetime of the object
		{ return total_sorts; }

	/// \param sorted_count [in/out] is the number of sorted elements in the beginning of neighbours array
	const neighbours & get(size_point p, const icolor & centre, uint & sorted_count);
	/// frees unused information during last pass, and marks remaining information as unused
	void pass_done();

public:
	const R & im;
	uint const whs;
	information const consider;
	uint const max_sorted;

private:
	typedef point<ushort> ushort_point; 
	struct cell {
		cell() { }
		cell(ushort_point r) : root(r), last_used(0) { }

		ushort_point root;
		neighbours nei;
		size_t sorted_count;
		size_t last_used;
	};
	typedef std::list<cell> precomputed_list;

private:
	cell & get_cell(size_point p);
	template <class D>
	void scan_neighbours(neighbours & n, size_point p, const icolor & centre, const D & distance);

private:
	img::simple_image<cell*> map; ///< only requested cells of the image contains valid data
	precomputed_list pool;
	uint total_scans, total_sorts;
};

template <class R>
sorted_neighbours<R>::sorted_neighbours(const R & im_, uint whs_, information consider_) :
	im(im_), whs(whs_), consider(consider_),
	max_sorted(std::max(7u, whs_ * (whs_ + 1))), //one third of pixels in circle window
	map(im_.dim())
{
	img::fill(map, (cell*)0);
	total_scans = total_sorts = 0;
}

template <class R>
typename sorted_neighbours<R>::cell & sorted_neighbours<R>::get_cell(size_point p)
{
	cell * & c = map[p];
	if (!c) {
		pool.push_back(cell(p));
		c = &pool.back();
	}
	else {
		c->last_used = 0;
		assert(c->root == p);
	}
	return *c;
}

template <class R>
void sorted_neighbours<R>::pass_done()
{
	precomputed_list::iterator i = pool.begin();
	while (i != pool.end()) {
		if (i->last_used < 2) {
			++i->last_used;
			++i;
		}
		else {
			map[size_point(i->root)] = 0;
			i = pool.erase(i);
		}
	}
}

struct color_distance {
	unsigned operator () (const icolor & delta) const
		{ return delta.sqr<unsigned>(); }
};

struct lightness_distance {
	unsigned operator () (const icolor & delta) const
		{ return abs(delta.r + delta.g + delta.b); }
};

struct hue_distance {
	unsigned operator () (const icolor & delta) const
		{ return (square(delta.r - delta.g) + square(delta.g - delta.b) + square(delta.b - delta.r))/2; }
};

template <class R>
template <class D>
void sorted_neighbours<R>::scan_neighbours(neighbours & n, size_point p, const icolor & centre, const D & distance)
{
	++total_scans;
	assert (whs <= std::numeric_limits<schar>::max());
	int_rect wnd;
	wnd.x1 = p.x >= whs ? -(int)whs : -(int)p.x;
    wnd.y1 = p.y >= whs ? -(int)whs : -(int)p.y;
	wnd.x2 = p.x + whs + 1 <= im.width() ? whs : im.width() - p.x;
	wnd.y2 = p.y + whs + 1 <= im.height() ? whs : im.height() - p.y;
	size_t area = (wnd.x2 - wnd.x1) * (wnd.y2 - wnd.y1);

	n.reserve(area - 1);

	spoint d;
	for (d.y = wnd.y1; d.y < wnd.y2; ++d.y) {
		for (d.x = wnd.x1; d.x < wnd.x2; ++d.x) 
		{
			icolor ref_col = im(p.x + d.x, p.y + d.y);
			if (d.x == 0 && d.y == 0 && centre == ref_col) //centre is equal to ref_col if they are taken from the same image, but may also be equal accidentally
				continue; //it's not right way for reference image
			if (whs > 2 && (d.x*d.x + d.y*d.y) > whs * whs)
				continue; //use circle window

			unsigned sim = distance(centre - ref_col);
			if (sim > std::numeric_limits<cdist_spoint_pair::first_type>::max())
				continue;

			n.push_back(cdist_spoint_pair(sim, d));
		}
	}

	assert (n.size() <= area - 1);
}

template <class R>
const neighbours & sorted_neighbours<R>::get(size_point p, const icolor & centre, uint & sorted_count)
{
	cell & c = get_cell(p);
	neighbours & n = c.nei;
	if (n.empty()) { 
		switch (consider) {
		case info_all:
			scan_neighbours(n, p, centre, color_distance());
			break;
		case info_lightness:
			scan_neighbours(n, p, centre, lightness_distance());
			break;
		case info_color:
			scan_neighbours(n, p, centre, hue_distance());
			break;
		default:
			assert(false);
		}
		c.sorted_count = 0;
	}
	else {
		if (c.sorted_count >= sorted_count || c.sorted_count == n.size()) {
			sorted_count = c.sorted_count;
			return n;
		}
	}

	uint max_sorted = std::min(this->max_sorted, n.size());

	if (sorted_count > max_sorted)
		sorted_count = max_sorted;
	std::partial_sort(n.begin(), n.begin() + sorted_count, n.end());
	++total_sorts;
	c.sorted_count = sorted_count;
	if (sorted_count == max_sorted && max_sorted < n.size()) {
		//remove tail of unsorted elements
		n.swap(neighbours(n.begin(), n.begin() + max_sorted));
	}

	return n;
}

} //namespace refine
